"""
Adam算法 （Kingma and Ba， 2014）

我们学习了：
随机梯度下降在解决优化问题时比梯度下降更有效
在一个小批量中使用更大的观测值集，可以通过向量化提供额外效率。
（高效的多机、多GPU和整体并行处理的关键）
添加一种机制，可以汇总过去梯度的历史加速收敛
对每个坐标缩放来实现高效计算机的预处理
通过学习率的调整来分离每个坐标的缩放。

Adam将上述所有技术都进行了汇总，是深度学习中更强大和有效的算法之一。
它的缺点：可能因为方差控制不良而发散，需要Yogi补丁来弥补这一问题
"""
import matplotlib.pyplot as plt
from d2l import torch as d2l
import torch


# TODO 有时间需要单独实现该算法的收敛
# TODO SGD moment
# TODO 了解学习率调度器

def init_adam_states(feature_dim):
    v_w, v_b = torch.zeros((feature_dim, 1)), torch.zeros(1)
    s_w, s_b = torch.zeros((feature_dim, 1)), torch.zeros(1)
    return ((v_w, s_w), (v_b, s_b))


def adam(params, states, hyperparams):
    beta1, beta2, eps = 0.9, 0.999, 1e-6
    for p, (v, s) in zip(params, states):
        with torch.no_grad():
            v[:] = beta1 * v + (1 - beta1) * p.grad
            s[:] = beta2 * s + (1 - beta2) * torch.square(p.grad)
            v_bias_corr = v / (1 - beta1 ** hyperparams['t'])
            s_bias_corr = s / (1 - beta2 ** hyperparams['t'])
            p[:] -= hyperparams['lr'] * v_bias_corr / (torch.sqrt(s_bias_corr)
                                                       + eps)
        p.grad.data.zero_()
    hyperparams['t'] += 1


data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
d2l.train_ch11(adam, init_adam_states(feature_dim),
               {'lr': 0.01, 't': 1}, data_iter, feature_dim)
plt.show()
